#include "arp_message.hh"
#include "ethernet_header.hh"
#include "ipv4_datagram.hh"
#include "network_interface_test_harness.hh"

#include <cstdlib>
#include <iostream>

using namespace std;

EthernetAddress random_private_ethernet_address() {
  EthernetAddress addr;
  for (auto &byte : addr) {
    byte = random_device()();  // use a random local Ethernet address
  }
  addr.at(0) |= 0x02;  // "10" in last two binary digits marks a private Ethernet address
  addr.at(0) &= 0xfe;

  return addr;
}

InternetDatagram make_datagram(const string &src_ip, const string &dst_ip) {
  InternetDatagram dgram;
  dgram.header().src = Address(src_ip, 0).ipv4_numeric();
  dgram.header().dst = Address(dst_ip, 0).ipv4_numeric();
  dgram.payload() = string("hello");
  dgram.header().len = dgram.header().hlen * 4 + dgram.payload().size();
  return dgram;
}

ARPMessage make_arp(const uint16_t opcode,
                    const EthernetAddress sender_ethernet_address,
                    const string sender_ip_address,
                    const EthernetAddress target_ethernet_address,
                    const string target_ip_address) {
  ARPMessage arp;
  arp.opcode = opcode;
  arp.sender_ethernet_address = sender_ethernet_address;
  arp.sender_ip_address = Address(sender_ip_address, 0).ipv4_numeric();
  arp.target_ethernet_address = target_ethernet_address;
  arp.target_ip_address = Address(target_ip_address, 0).ipv4_numeric();
  return arp;
}

EthernetFrame make_frame(const EthernetAddress &src,
                         const EthernetAddress &dst,
                         const uint16_t type,
                         const BufferList payload) {
  EthernetFrame frame;
  frame.header().src = src;
  frame.header().dst = dst;
  frame.header().type = type;
  frame.payload() = payload.concatenate();
  return frame;
}

int main() {
  try {
    {
      const EthernetAddress local_eth = random_private_ethernet_address();
      NetworkInterfaceTestHarness test{"typical ARP workflow", local_eth, Address("4.3.2.1", 0)};

      const auto datagram = make_datagram("5.6.7.8", "13.12.11.10");
      test.execute(SendDatagram{datagram, Address("192.168.0.1", 0)});

      // outgoing datagram should result in an ARP request
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, local_eth, "4.3.2.1", {}, "192.168.0.1").serialize())});
      test.execute(ExpectNoFrame{});
      const EthernetAddress target_eth = random_private_ethernet_address();

      test.execute(Tick{800});
      test.execute(ExpectNoFrame{});

      // ARP reply should result in the queued datagram getting sent
      test.execute(ReceiveFrame{
          make_frame(target_eth,
                     local_eth,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REPLY, target_eth, "192.168.0.1", local_eth, "4.3.2.1").serialize()),
          {}});

      test.execute(ExpectFrame{make_frame(local_eth, target_eth, EthernetHeader::TYPE_IPv4, datagram.serialize())});
      test.execute(ExpectNoFrame{});

      // any IP reply directed for our Ethernet address should be passed up the stack
      const auto reply_datagram = make_datagram("13.12.11.10", "5.6.7.8");
      test.execute(ReceiveFrame(
          make_frame(target_eth, local_eth, EthernetHeader::TYPE_IPv4, reply_datagram.serialize()), reply_datagram));
      test.execute(ExpectNoFrame{});

      // incoming frames to another Ethernet address (not ours) should be ignored
      const EthernetAddress another_eth = {1, 1, 1, 1, 1, 1};
      test.execute(
          ReceiveFrame(make_frame(target_eth, another_eth, EthernetHeader::TYPE_IPv4, reply_datagram.serialize()), {}));
    }

    {
      const EthernetAddress local_eth = random_private_ethernet_address();
      const EthernetAddress remote_eth = random_private_ethernet_address();
      NetworkInterfaceTestHarness test{"reply to ARP request", local_eth, Address("5.5.5.5", 0)};
      test.execute(ReceiveFrame{
          make_frame(remote_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, remote_eth, "10.0.1.1", {}, "7.7.7.7").serialize()),
          {}});
      test.execute(ExpectNoFrame{});
      test.execute(ReceiveFrame{
          make_frame(remote_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, remote_eth, "10.0.1.1", {}, "5.5.5.5").serialize()),
          {}});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     remote_eth,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REPLY, local_eth, "5.5.5.5", remote_eth, "10.0.1.1").serialize())});
      test.execute(ExpectNoFrame{});
    }

    {
      const EthernetAddress local_eth = random_private_ethernet_address();
      const EthernetAddress remote_eth = random_private_ethernet_address();
      NetworkInterfaceTestHarness test{"learn from ARP request", local_eth, Address("5.5.5.5", 0)};
      test.execute(ReceiveFrame{
          make_frame(remote_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, remote_eth, "10.0.1.1", {}, "5.5.5.5").serialize()),
          {}});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     remote_eth,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REPLY, local_eth, "5.5.5.5", remote_eth, "10.0.1.1").serialize())});
      test.execute(ExpectNoFrame{});

      const auto datagram = make_datagram("5.6.7.8", "13.12.11.10");
      test.execute(SendDatagram{datagram, Address("10.0.1.1", 0)});
      test.execute(ExpectFrame{make_frame(local_eth, remote_eth, EthernetHeader::TYPE_IPv4, datagram.serialize())});
      test.execute(ExpectNoFrame{});
    }

    {
      const EthernetAddress local_eth = random_private_ethernet_address();
      NetworkInterfaceTestHarness test{"pending mappings last five seconds", local_eth, Address("1.2.3.4", 0)};

      test.execute(SendDatagram{make_datagram("5.6.7.8", "13.12.11.10"), Address("10.0.0.1", 0)});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, local_eth, "1.2.3.4", {}, "10.0.0.1").serialize())});
      test.execute(ExpectNoFrame{});
      test.execute(Tick{4990});
      test.execute(SendDatagram{make_datagram("17.17.17.17", "18.18.18.18"), Address("10.0.0.1", 0)});
      test.execute(ExpectNoFrame{});
      test.execute(Tick{20});
      // pending mapping should now expire
      test.execute(SendDatagram{make_datagram("42.41.40.39", "13.12.11.10"), Address("10.0.0.1", 0)});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, local_eth, "1.2.3.4", {}, "10.0.0.1").serialize())});
      test.execute(ExpectNoFrame{});
    }

    {
      const EthernetAddress local_eth = random_private_ethernet_address();
      NetworkInterfaceTestHarness test{"active mappings last 30 seconds", local_eth, Address("4.3.2.1", 0)};

      const auto datagram = make_datagram("5.6.7.8", "13.12.11.10");
      const auto datagram2 = make_datagram("5.6.7.8", "13.12.11.11");
      const auto datagram3 = make_datagram("5.6.7.8", "13.12.11.12");
      const auto datagram4 = make_datagram("5.6.7.8", "13.12.11.13");

      test.execute(SendDatagram{datagram, Address("192.168.0.1", 0)});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, local_eth, "4.3.2.1", {}, "192.168.0.1").serialize())});
      const EthernetAddress target_eth = random_private_ethernet_address();
      test.execute(ReceiveFrame{
          make_frame(target_eth,
                     local_eth,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REPLY, target_eth, "192.168.0.1", local_eth, "4.3.2.1").serialize()),
          {}});

      test.execute(ExpectFrame{make_frame(local_eth, target_eth, EthernetHeader::TYPE_IPv4, datagram.serialize())});
      test.execute(ExpectNoFrame{});

      // try after 10 seconds -- no ARP should be generated
      test.execute(Tick{10000});
      test.execute(SendDatagram{datagram2, Address("192.168.0.1", 0)});
      test.execute(ExpectFrame{make_frame(local_eth, target_eth, EthernetHeader::TYPE_IPv4, datagram2.serialize())});
      test.execute(ExpectNoFrame{});

      // another 10 seconds -- no ARP should be generated
      test.execute(Tick{10000});
      test.execute(SendDatagram{datagram3, Address("192.168.0.1", 0)});
      test.execute(ExpectFrame{make_frame(local_eth, target_eth, EthernetHeader::TYPE_IPv4, datagram3.serialize())});
      test.execute(ExpectNoFrame{});

      // after another 11 seconds, need to ARP again
      test.execute(Tick{11000});
      test.execute(SendDatagram{datagram4, Address("192.168.0.1", 0)});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, local_eth, "4.3.2.1", {}, "192.168.0.1").serialize())});
      test.execute(ExpectNoFrame{});

      // ARP reply again
      const EthernetAddress new_target_eth = random_private_ethernet_address();
      test.execute(ReceiveFrame{
          make_frame(
              new_target_eth,
              local_eth,
              EthernetHeader::TYPE_ARP,
              make_arp(ARPMessage::OPCODE_REPLY, new_target_eth, "192.168.0.1", local_eth, "4.3.2.1").serialize()),
          {}});
      test.execute(
          ExpectFrame{make_frame(local_eth, new_target_eth, EthernetHeader::TYPE_IPv4, datagram4.serialize())});
      test.execute(ExpectNoFrame{});
    }

    {
      const EthernetAddress local_eth = random_private_ethernet_address();
      const EthernetAddress remote_eth1 = random_private_ethernet_address();
      const EthernetAddress remote_eth2 = random_private_ethernet_address();
      NetworkInterfaceTestHarness test{"different ARP mappings are independent", local_eth, Address("10.0.0.1", 0)};

      // first ARP mapping
      test.execute(ReceiveFrame{
          make_frame(remote_eth1,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, remote_eth1, "10.0.0.5", {}, "10.0.0.1").serialize()),
          {}});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     remote_eth1,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REPLY, local_eth, "10.0.0.1", remote_eth1, "10.0.0.5").serialize())});
      test.execute(ExpectNoFrame{});

      test.execute(Tick{15000});

      // second ARP mapping
      test.execute(ReceiveFrame{
          make_frame(remote_eth2,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, remote_eth2, "10.0.0.19", {}, "10.0.0.1").serialize()),
          {}});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     remote_eth2,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REPLY, local_eth, "10.0.0.1", remote_eth2, "10.0.0.19").serialize())});
      test.execute(ExpectNoFrame{});

      test.execute(Tick{10000});

      // outgoing datagram to first destination
      const auto datagram = make_datagram("5.6.7.8", "13.12.11.10");
      test.execute(SendDatagram{datagram, Address("10.0.0.5", 0)});

      // outgoing datagram to second destination
      const auto datagram2 = make_datagram("100.99.98.97", "4.10.4.10");
      test.execute(SendDatagram{datagram2, Address("10.0.0.19", 0)});

      test.execute(ExpectFrame{make_frame(local_eth, remote_eth1, EthernetHeader::TYPE_IPv4, datagram.serialize())});
      test.execute(ExpectFrame{make_frame(local_eth, remote_eth2, EthernetHeader::TYPE_IPv4, datagram2.serialize())});
      test.execute(ExpectNoFrame{});

      test.execute(Tick{5010});

      // outgoing datagram to second destination (mapping still alive)
      const auto datagram3 = make_datagram("150.140.130.120", "144.144.144.144");
      test.execute(SendDatagram{datagram3, Address("10.0.0.19", 0)});
      test.execute(ExpectFrame{make_frame(local_eth, remote_eth2, EthernetHeader::TYPE_IPv4, datagram3.serialize())});
      test.execute(ExpectNoFrame{});

      // outgoing datagram to second destination (mapping has expired)
      const auto datagram4 = make_datagram("244.244.244.244", "3.3.3.3");
      test.execute(SendDatagram{datagram4, Address("10.0.0.5", 0)});
      test.execute(ExpectFrame{
          make_frame(local_eth,
                     ETHERNET_BROADCAST,
                     EthernetHeader::TYPE_ARP,
                     make_arp(ARPMessage::OPCODE_REQUEST, local_eth, "10.0.0.1", {}, "10.0.0.5").serialize())});
      test.execute(ExpectNoFrame{});
    }
  } catch (const exception &e) {
    cerr << e.what() << endl;
    return EXIT_FAILURE;
  }

  return EXIT_SUCCESS;
}
